"""
模型加载器模块
负责管理和加载各种深度学习模型
"""

import os
import requests
from pathlib import Path
from typing import Optional, Dict, Any
from tqdm import tqdm
import hashlib


class ModelLoader:
    """
    模型加载器类
    
    功能：
    - 自动下载缺失的模型
    - 验证模型完整性
    - 管理模型版本
    """
    
    def __init__(self, models_dir: str = "models", config: Optional[Dict] = None):
        """
        初始化模型加载器
        
        参数:
            models_dir (str): 模型存储目录
            config (Dict, optional): 配置字典，包含下载地址等信息
        """
        self.models_dir = Path(models_dir)
        self.models_dir.mkdir(parents=True, exist_ok=True)
        
        # 从配置文件加载模型信息
        if config and 'model_download' in config:
            self.model_configs = config['model_download']['urls']
            self.download_settings = {
                'timeout': config['model_download'].get('timeout', 30),
                'retry_count': config['model_download'].get('retry_count', 3),
                'chunk_size': config['model_download'].get('chunk_size', 8192)
            }
        else:
            # 默认配置（如果没有配置文件）
            self.model_configs = {
                'yolov5_plate': {
                    'url': 'https://github.com/ultralytics/assets/releases/download/v8.3.0/yolov5nu.pt',
                    'backup_urls': ['https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5n.pt'],
                    'md5': None,
                    'size': '5.3MB',
                    'description': 'YOLOv5车牌检测模型（兼容版本）'
                }
            }
            self.download_settings = {
                'timeout': 30,
                'retry_count': 3,
                'chunk_size': 8192
            }
        
    def load_model(self, model_name: str, 
                   force_download: bool = False) -> Optional[Path]:
        """
        加载指定模型
        
        参数:
            model_name (str): 模型名称
            force_download (bool): 是否强制重新下载
            
        返回:
            Optional[Path]: 模型文件路径，失败返回None
        """
        if model_name not in self.model_configs:
            print(f"未知的模型: {model_name}")
            return None
            
        model_config = self.model_configs[model_name]
        model_path = self.models_dir / f"{model_name}.pt"
        
        # 检查模型是否存在
        if model_path.exists() and not force_download:
            # 验证模型完整性
            if self._verify_model(model_path, model_config.get('md5')):
                print(f"模型 {model_name} 已存在且完整")
                return model_path
            else:
                print(f"模型 {model_name} 损坏，重新下载")
                
        # 下载模型（尝试主URL和备用URL）
        urls_to_try = [model_config['url']]
        if 'backup_urls' in model_config:
            urls_to_try.extend(model_config['backup_urls'])
            
        for url in urls_to_try:
            if self._download_model(model_name, url, model_path):
                return model_path
                
        return None
            
    def _download_model(self, model_name: str, 
                       url: str, 
                       save_path: Path) -> bool:
        """
        下载模型文件
        
        参数:
            model_name (str): 模型名称
            url (str): 下载URL
            save_path (Path): 保存路径
            
        返回:
            bool: 是否下载成功
        """
        retry_count = self.download_settings['retry_count']
        
        for attempt in range(retry_count):
            try:
                print(f"正在下载 {model_name} 从 {url} (尝试 {attempt + 1}/{retry_count})")
                
                # 发送请求
                response = requests.get(
                    url, 
                    stream=True, 
                    timeout=self.download_settings['timeout']
                )
                response.raise_for_status()
                
                # 获取文件大小
                total_size = int(response.headers.get('content-length', 0))
                
                # 下载并显示进度
                with open(save_path, 'wb') as f:
                    with tqdm(total=total_size, unit='B', 
                             unit_scale=True, desc=model_name) as pbar:
                        for chunk in response.iter_content(
                            chunk_size=self.download_settings['chunk_size']
                        ):
                            if chunk:
                                f.write(chunk)
                                pbar.update(len(chunk))
                                
                print(f"{model_name} 下载完成")
                return True
                
            except Exception as e:
                print(f"下载 {model_name} 失败 (尝试 {attempt + 1}): {e}")
                if save_path.exists():
                    save_path.unlink()  # 删除不完整的文件
                    
                if attempt < retry_count - 1:
                    print(f"等待重试...")
                    import time
                    time.sleep(2)  # 等待2秒后重试
                    
        return False
            
    def _verify_model(self, model_path: Path, 
                     expected_md5: Optional[str]) -> bool:
        """
        验证模型文件的完整性
        
        参数:
            model_path (Path): 模型文件路径
            expected_md5 (str, optional): 期望的MD5值
            
        返回:
            bool: 是否验证通过
        """
        if not expected_md5:
            # 如果没有提供MD5，只检查文件是否存在且非空
            return model_path.exists() and model_path.stat().st_size > 0
            
        # 计算文件MD5
        md5_hash = hashlib.md5()
        with open(model_path, 'rb') as f:
            for chunk in iter(lambda: f.read(4096), b""):
                md5_hash.update(chunk)
                
        actual_md5 = md5_hash.hexdigest()
        return actual_md5 == expected_md5
        
    def list_models(self) -> Dict[str, Any]:
        """
        列出所有可用的模型
        
        返回:
            Dict[str, Any]: 模型信息字典
        """
        models_info = {}
        
        for model_name, config in self.model_configs.items():
            model_path = self.models_dir / f"{model_name}.pt"
            models_info[model_name] = {
                'description': config['description'],
                'size': config['size'],
                'installed': model_path.exists(),
                'path': str(model_path) if model_path.exists() else None,
                'url': config['url'],
                'backup_urls': config.get('backup_urls', [])
            }
            
        return models_info
        
    def clean_models(self) -> None:
        """
        清理所有下载的模型文件
        """
        for model_file in self.models_dir.glob("*"):
            if model_file.is_file():
                model_file.unlink()
                print(f"已删除: {model_file.name}")
                
    def get_model_path(self, model_name: str) -> Optional[Path]:
        """
        获取模型文件路径
        
        参数:
            model_name (str): 模型名称
            
        返回:
            Optional[Path]: 模型路径，不存在返回None
        """
        model_path = self.models_dir / f"{model_name}.pt"
        return model_path if model_path.exists() else None 